#!/usr/bin/env python
# s3funnel - Multithreaded tool for performing operations on Amazon's S3
# Copyright (c) 2008 Andrey Petrov
#
# This module is part of s3funnel and is released under
# the MIT license: http://www.opensource.org/licenses/mit-license.php

"""
s3funnel is a multithreaded tool for performing operations on Amazon's S3.

Key Operations:
    DELETE Delete key from the bucket
    GET    Get key from the bucket
    PUT    Put file into the bucket (key is the basename of the path)

Bucket Operations:
    CREATE Create a new bucket
    DROP   Delete an existing bucket (must be empty)
    LIST   List keys in the bucket. If no bucket is given, buckets will be listed.
"""

__VERSION__ = '0.6'

"""
TODO:
* Use callback to do an interactive progress bar
"""

import logging

log = logging.getLogger(__name__)
s3funnel_log = logging.getLogger('s3funnel')
boto_log = logging.getLogger('boto')

def set_log_level(level):
    for i in [log, s3funnel_log, boto_log]:
        i.setLevel(level)


# Official Python modules
import os
import sys
import signal
import threading
import socket
import httplib

from glob import glob
from optparse import OptionParser, SUPPRESS_HELP

# Third party modules
import workerpool

# Local imports
from s3funnel import S3Funnel
from s3funnel.exceptions import FunnelError

event_stop = threading.Event()

def main():
    # Parse the command line...
    usage="%prog BUCKET OPERATION [OPTIONS] [FILE]...\n" + __doc__
    parser = OptionParser(usage)
    parser.add_option("-a", "--aws_key",    dest="aws_key", type="string", help="Overrides AWS_ACCESS_KEY_ID environment variable")
    parser.add_option("-s", "--aws_secret_key", dest="aws_secret_key", type="string", help="Overrides AWS_SECRET_ACCESS_KEY environment variable")
    parser.add_option("-t", "--threads",    dest="numthreads", default=1, type="int", metavar="N", help="Number of threads to use [default: %default]")
    parser.add_option("-T", "--timeout",    dest="timeout", default=0, type="float", metavar="SECONDS", help="Socket timeout time, 0 is never [default: %default]")
    parser.add_option("--insecure",         dest="secure", action="store_false", default=True, help="Don't use secure (https) connection")
    parser.add_option("--list-marker",      dest="list_marker", type="string", default=None, metavar="KEY", help="(`list` only) Start key for list operation")
    parser.add_option("--list-prefix",      dest="list_prefix", type="string", default=None, metavar="STRING", help="(`list` only) Limit results to a specific prefix")
    parser.add_option("--list-delimiter",   dest="list_delimiter", type="string", default=None, metavar="CHAR", help="(`list` only) Treat value as a delimiter for hierarchical listing")
    parser.add_option("--put-acl",          dest="acl", type="string", default="public-read", help="(`put` only) Set the ACL permission for each file [default: %default]")
    parser.add_option("--put-full-path",    dest="put_full_path", action="store_true", help="(`put` only) Use the full given path as the key name, instead of just the basename")
    parser.add_option("--put-only-new",     dest="put_only_new", action="store_true", help="(`put` only) Only PUT keys which don't already exist in the bucket with the same md5 digest")
    parser.add_option("--put-header",       dest="headers", type="string", action="append", help="(`put` only) Add the specified header to the request")
    parser.add_option("-i", "--input",      dest="input", type="string", metavar="FILE", help="Read one file per line from a FILE manifest")
    parser.add_option("-v", "--verbose",    dest="verbose", action="count", default=None, help="Enable verbose output. Use twice to enable debug output")
    parser.add_option("--version",          dest="version", action="store_true", help="Output version information and exit")

    # Deprecated options (for backwards compatibility)
    parser.add_option("--start_key",        dest="list_marker", type="string", default=None, help=SUPPRESS_HELP)
    parser.add_option("--acl",              dest="acl", type="string", default="public-read", help=SUPPRESS_HELP)

    # Hidden options for debugging, development, and lols
    # TODO: ...

    options, args = parser.parse_args()

    # Version?
    if options.version:
        print "s3funnel %s" % __VERSION__
        return 0

    # Check input
    aws_key = os.environ.get("AWS_ACCESS_KEY_ID")
    aws_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY")

    ## AWS
    if options.aws_key:
        aws_key = options.aws_key
    if options.aws_secret_key:
        aws_secret_key = options.aws_secret_key
    if None in [aws_key, aws_secret_key]:
        parser.error("Missing required arguments `aws_key` or `aws_secret_key`")

    ## Threads
    if options.numthreads < 1:
        parser.error("`theads` must be at least 1")

    ## Misc. options
    if options.timeout:
        try:
            socket.setdefaulttimeout(options.timeout)
        except TypeError, e:
            parser.error("`timeout` error: %s" % e.message)

    ## Parse put headers
    headers = {}
    if options.headers:
        for header in options.headers:
            if ':' not in header:
                parser.error("Header must be use ':' separator")
            key, value = header.split(':', 2)
            headers[key.strip()] = value.strip()

    # Arguments
    if len(args) < 1:
        parser.error("BUCKET not specified")
    bucket = args[0]
    if len(args) < 2:
        # Exception for single-argument operations
        if bucket.lower() in ['list']:
            operation = bucket.lower()
            bucket = None
        else:
            parser.error("OPERATION not specified")
    else:
        operation = args[1].lower()

    if operation == 'list' and not bucket:
        operation = 'show'

    # Setup logging
    if options.verbose > 1:
        set_log_level(logging.DEBUG)
    elif options.verbose > 0:
        set_log_level(logging.INFO)

    # Setup operation configuration
    config = {'acl': options.acl,
              'list_marker': options.list_marker or '',
              'list_prefix': options.list_prefix or '',
              'list_delimiter': options.list_delimiter or '',
              'aws_key': aws_key,
              'aws_secret_key': aws_secret_key,
              'secure': options.secure,
              'bucket': bucket,
              'put_full_path': options.put_full_path,
              'put_only_new': options.put_only_new,
              'headers': headers,
              'numthreads': options.numthreads
              }

    funnel = S3Funnel(**config)

    # Setup operation mapping
    methods_keys = {
       'get':    funnel.get,
       'put':    funnel.put,
       'delete': funnel.delete,
    }

    methods_bucket = {
       'list':   funnel.list_bucket,
       'drop':   funnel.drop_bucket,
       'create': funnel.create_bucket,
    }

    methods_global = {
       'show':   funnel.show_buckets,
    }

    valid_operations = methods_keys.keys() + methods_bucket.keys() + methods_global.keys()
    valid_operations.sort()
    if operation not in valid_operations:
        parser.error("OPERATION must be one of: %s" % ', '.join(valid_operations))

    # Get data source
    input_src = None
    if options.input:
        # Get source from manifest or stdin (via -i flag)
        if options.input == '-':
            input_src = "stdin"
            options.input = sys.stdin
        try:
            data = open(glob(options.input)[0])
        except (IOError, IndexError), e:
            log.error("%s: File not found" % options.input)
            return -1
        input_src = "`%s'" % options.input
    elif len(args) < 3:
        # Get source from stdin
        input_src = "stdin"
        data = sys.stdin
    else:
        if operation == 'put':
            # Get source from glob-expanded arguments
            data = []
            for arg in args[2:]:
                found = glob(arg)
                if not found:
                    log.error("%s: No such file." % arg)
                    continue
                data += found
        else:
            data = args[2:]
        input_src = "arguments: %s" % ', '.join(data)

    # Setup interrupt handling
    def shutdown(signum, stack):
        log.warning("Interrupted, shutting down...")
        event_stop.set()
    signal.signal(signal.SIGINT, shutdown)

    # Setup output logger
    output = logging.getLogger('output')
    output_handler = logging.StreamHandler(sys.stdout)
    output_handler.setFormatter(logging.Formatter('%(message)s'))
    output.addHandler(output_handler)
    output.setLevel(logging.INFO)

    # Feed input into the appropriate method
    log.info("Using input from %s" % input_src)

    try:
        # TODO: Rewrite this into something fancier
        if operation in methods_global:
            m = methods_global[operation]
            for i in m():
                output.info(i)
        elif operation in methods_bucket:
            m = methods_bucket[operation]
            for i in m(bucket, **config):
                output.info(i)
        elif operation in methods_keys:
            m = methods_keys[operation]
            r = m(ikeys=(d.strip() for d in data), **config)
            if r:
                log.critical("%d keys failed." % len(r))
                return len(r)
    except FunnelError, e:
        log.critical(e.message)
        return -1

    funnel.shutdown()

if __name__ == "__main__":
    log_handler = logging.StreamHandler()
    log_handler.setFormatter(logging.Formatter('%(levelname)-8s %(message)s'))

    log.addHandler(log_handler)
    s3funnel_log.addHandler(log_handler)
    boto_log.addHandler(log_handler)

    r = main()
    if r:
        sys.exit(r)
